#include "key_switcher.h"
/*
 * @brief 密钥切换器的构造函数
 * @param params 上下文参数指针，包含环的参数和其他全局设置
 */
KeySwitcher::KeySwitcher(Context *params) : params(params)
{
    baseconverter = new FastBasisExtender(params->ringQ, params->ringP);

    int beta = params->Beta();

    pool.resize(6);

    //test
    poolTest0.resize(16);
    poolTest1.resize(16);
    for (int i = 0; i < 16; i++) {
        poolTest0[i] = params->NewRingQP()->NewPoly();
        poolTest1[i] = params->NewRingQP()->NewPoly();
    }

    for (int i = 0; i < 6; i++) {
        pool[i] = params->NewRingQP()->NewPoly();
    }

    poolInvNTT = params->ringQ->NewPoly();

    poolDecompQP.resize(beta);
    for (int i = 0; i < beta; i++) {
        poolDecompQP[i] = params->NewRingQP()->NewPoly();
    }
}
/*
 * @brief 对密文执行密钥切换操作
 *
 * 在给定层级上执行密钥切换,包含两个步骤:
 * 1. 调用SwitchKeysInPlaceNoModDown进行密钥切换但不执行模下降
 * 2. 执行模下降操作将结果转换回原始模数空间
 *
 * @param levelQ 操作的Q基层级
 * @param cx 输入密文
 * @param evakey 切换密钥
 * @param p0 输出密文的第一个多项式
 * @param p1 输出密文的第二个多项式
 */
void KeySwitcher::SwitchKeysInPlace(int levelQ, Poly *cx, SwitchingKey *evakey, Poly *p0, Poly *p1)
{
    SwitchKeysInPlaceNoModDown(levelQ, cx, evakey, p0, pool[1]->P, p1, pool[2]->P);
    uint64_t levelP = evakey->value[0][0]->P->Coeffs.size() - 1;

    if (cx->isNTT) {
        params->ringP->InvNTTLvl(levelP, pool[1]->P, pool[1]->P);
        params->ringP->InvNTTLvl(levelP, pool[2]->P, pool[1]->P);
        baseconverter->BaseConvertQP2Q(levelQ, levelP, p0, pool[1]->P, p0);
        baseconverter->BaseConvertQP2Q(levelQ, levelP, p1, pool[1]->P, p1);
        params->ringQ->NTTLvl(levelQ, p0, p0);
        params->ringQ->NTTLvl(levelQ, p1, p1);
    } else {
        params->ringQ->InvNTTLvl(levelQ, p0, p0);
        params->ringQ->InvNTTLvl(levelQ, p1, p1);
        params->ringP->InvNTTLvl(levelP, pool[1]->P, pool[1]->P);
        params->ringP->InvNTTLvl(levelP, pool[2]->P, pool[2]->P);
        baseconverter->BaseConvertQP2Q(levelQ, levelP, p0, pool[1]->P, p0);
        baseconverter->BaseConvertQP2Q(levelQ, levelP, p1, pool[2]->P, p1);
    }
}
/*
 * @brief 执行密钥切换的核心操作,不包含模下降步骤
 *
 * @param levelQ 操作的Q基层级
 * @param cx 输入密文
 * @param evakey 切换密钥
 * @param c0Q 输出密文在Q基的第一个多项式
 * @param c0P 输出密文在P基的第一个多项式
 * @param c1Q 输出密文在Q基的第二个多项式
 * @param c1P 输出密文在P基的第二个多项式
 */
void KeySwitcher::SwitchKeysInPlaceNoModDown(int levelQ, Poly *cx, SwitchingKey *evakey, Poly *c0Q, Poly *c0P,
    Poly *c1Q, Poly *c1P)
{
    int reduce;

    Ring *ringQ = params->ringQ;
    Ring *ringP = params->ringP;
    RingQP *ringQP = params->NewRingQP();

    PolyQP *c2QP = pool[0];

    Poly *cxNTT, *cxInvNTT;
    if (cx->isNTT) {
        cxNTT = cx;
        cxInvNTT = poolInvNTT;
        ringQ->InvNTTLvl(levelQ, cxNTT, cxInvNTT);
    } else {
        cxNTT = poolInvNTT;
        cxInvNTT = cx;
        ringQ->NTTLvl(levelQ, cxInvNTT, cxNTT);
    }

    PolyQP *c0QP = new PolyQP(c0Q, c0P);
    PolyQP *c1QP = new PolyQP(c1Q, c1P);

    reduce = 0;

    uint64_t alpha = evakey->value[0][0]->P->Coeffs.size();
    uint64_t levelP = alpha - 1;
    int beta = int(ceil(double(levelQ + 1) / double(levelP + 1)));

    int QiOverF = params->QiOverflowMargin(levelQ) >> 1;

    // Key switching with CRT decomposition for the Qi

    #pragma omp parrallel for
    for (int i = 0; i < beta; i++) {
        DecomposeSingleNTT(levelQ, levelP, alpha, i, cxNTT, cxInvNTT, c2QP->Q, c2QP->P);
        if (i == 0) {
            ringQP->MulCoeffsBarrettConstantLvl(levelQ, levelP, evakey->value[i][0], c2QP, c0QP);
            ringQP->MulCoeffsBarrettConstantLvl(levelQ, levelP, evakey->value[i][1], c2QP, c1QP);
        } else {
            ringQP->MulCoeffsBarrettConstantAndAddNoModLvl(levelQ, levelP, evakey->value[i][0], c2QP, c0QP);
            ringQP->MulCoeffsBarrettConstantAndAddNoModLvl(levelQ, levelP, evakey->value[i][1], c2QP, c1QP);
        }
        reduce++;
    }

    if ((reduce % QiOverF) != 0) {
        #pragma omp parrallel for collapse(2)
        for (int i = 0; i < levelQ + 1; i++) {
            for (uint64_t j = 0; j < ringQ->N; j++) {
                c0QP->Q->Coeffs[i][j] = BRedAdd(c0QP->Q->Coeffs[i][j], ringQ->Modulus[i], ringQ->bredParams[i]);
                c1QP->Q->Coeffs[i][j] = BRedAdd(c1QP->Q->Coeffs[i][j], ringQ->Modulus[i], ringQ->bredParams[i]);
            }
        }
        for (int i = 0; i < levelP + 1; i++) {
            #pragma omp parrallel for 
            for (uint64_t j = 0; j < ringP->N; j++) {
                c0QP->P->Coeffs[i][j] = BRedAdd(c0QP->P->Coeffs[i][j], ringP->Modulus[i], ringP->bredParams[i]);
                c1QP->P->Coeffs[i][j] = BRedAdd(c1QP->P->Coeffs[i][j], ringP->Modulus[i], ringP->bredParams[i]);
            }
        }
    }
}

/*
 * @brief 对NTT域中的单个多项式进行分解
 *
 * 将输入多项式在NTT域中分解为多个小系数多项式:
 * 1. 对输入执行逆NTT变换
 * 2. 按照alpha和beta参数进行基分解
 * 3. 对分解结果执行NTT变换
 *
 * @param levelQ Q基的层级
 * @param levelP P基的层级
 * @param alpha 分解的块大小
 * @param beta 分解的块数量
 * @param c2NTT NTT域中的输入多项式
 * @param c2InvNTT 用于存储逆NTT结果的临时多项式
 * @param c2QiQ 分解结果在Q基中的系数
 * @param c2QiP 分解结果在P基中的系数
 */
void KeySwitcher::DecomposeSingleNTT(int levelQ, int levelP, int alpha, int beta, Poly *c2NTT, Poly *c2InvNTT,
    Poly *c2QiQ, Poly *c2QiP)
{
    Ring *ringQ = params->ringQ;
    Ring *ringP = params->ringP;

    decomposer->DecomposeAndSplit(levelQ, levelP, alpha, beta, c2InvNTT, c2QiQ, c2QiP);

    int p0idxst = beta * (levelP + 1);
    int p0idxed = p0idxst + 1;
    ringQ->NTT(c2QiQ, c2QiQ);
    ringP->NTT(c2QiP, c2QiP);
}